import numpy as np
import itertools


class IndefinitEquation:
    '''
    4x_1 + 5x_2 - x_3 + x_4 + 6x_5 = 3
    3x_1 - 9x_2 - 7x_3 + 5x_4 - 10x_5 = 1
    5x_1 + 2x_2 + x_3 - 4x_4 + 7x_5 = 2
    对于单个不定方程，构造
    A = [[a1  1  0  ......  0]
         [a2  0  1  ......  0]
         [...  ..............]
         [an  0  0  ......  1]]
    变换得到
    B = [[r  b11  b21  ......  bn1]
         [0  b12  b22  ......  bn2]
         [...  ........     ......]
         [0  b1n  b2n  ......  bnn]]
    |r|是a1,a2...an的最大公因数，所有b均为整数
    '''
    def __init__(self, a, target):
        self.target = np.array(target).reshape(-1, 1)
        self.row_target = self.target.shape[0]   # 通过目标值得到方程个数
        self.a = np.array(a).reshape(self.row_target, -1)    # 调整系数shape
        self.cols = self.a.shape[1]
        self.A = np.c_[self.a.T, np.identity(self.cols)]      # 右接对角单位矩阵

    # 计算最大公约数
    def gcd(self, num1, num2):
        while num1 != num2:
            if 0 in [num1, num2]:
                return abs([num1, num2].remove(0)[0])
            if num1 < num2:
                if num2 < 2*num1:
                    num2 -= num1
                else:
                    num2 -= int(num2/num1-1)*num1
            else:
                if num1 < 2*num2:
                    num1 -= num2
                else:
                    num1 -= int(num1/num2)*num2
        return num1

    # 封装计算多个数字最大公约数
    def pack_gcd(self, num_ls):
        num1 = abs(num_ls.pop())
        while num_ls:
            num2 = abs(num_ls.pop())
            num1 = self.gcd(num1, num2)
        return int(num1)

    # 初等变换将A变为B
    def change_A_to_B(self):
        # 将左下角先归0
        for col in range(self.cols):
            if col < self.row_target:
                i = col
                while list(self.A[col:, col]).count(0) < self.cols - col - 1:
                    if i == self.cols - 1:  # i增加到最底部时重置到对角线处
                        i = col
                    if self.A[i, col] == 0:  # 当前位置为0则跳到下一位进行新一轮循环
                        i += 1
                        continue
                    j = 1  # 非零数字索引调整
                    try:
                        while self.A[i + j, col] == 0 and i + j <= self.cols - 1:
                            j += 1
                    except IndexError:  # 索引加到最底部都是0则将i重置到对角线处
                        i = col
                        continue
                    if abs(self.A[i, col]) == abs(self.A[i + j, col]):  # 先比较绝对值大小，再分正负
                        if self.A[i, col] * self.A[i + j, col] >= 0:
                            self.A[i + j] -= self.A[i]
                        else:
                            self.A[i + j] += self.A[i]
                    elif abs(self.A[i, col]) < abs(self.A[i + j, col]):
                        self.A[i + j] -= int(self.A[i + j, col] / self.A[i, col]) * self.A[i]
                    else:
                        if i > col:
                            self.A[i] -= int(self.A[i, col] / self.A[i + j, col]) * self.A[i + j]
                        else:  # 此时只有i:=col
                            if int(self.A[i, col] / self.A[i + j, col]) >= 2:  # 2倍及以上，相减时减少1倍
                                self.A[i] -= int(self.A[i, col] / self.A[i + j, col] - 1) * self.A[i + j]
                            elif int(self.A[i, col] / self.A[i + j, col]) <= -2:
                                self.A[i] -= int(self.A[i, col] / self.A[i + j, col] + 1) * self.A[i + j]
                            else:
                                self.A[i] -= int(self.A[i, col] / self.A[i + j, col]) * self.A[i + j]
                    i += 1
            else:
                i = 0
                gcd = self.pack_gcd(list(self.A[col:, col]))
                while list(map(abs, self.A[col:, col])).count(gcd) < self.cols - col:
                    if i == self.cols - 1:     # i增加到最底部时重置到对角线处
                        i = 0
                    if self.A[i, col] == gcd:  # 当前位置等于公约数则跳到下一位进行新一轮循环
                        i += 1
                        continue
                    j = 1  # 非公约数索引调整
                    try:
                        self.A[i + j]
                    except IndexError:
                        i = 0
                        continue
                    # 此时应考虑不能归0处理
                    if abs(self.A[i, col]) > abs(self.A[i + j, col]) and self.A[i + j, col] != 0:
                        if self.A[i, col] / self.A[i + j, col] == self.A[i, col] // self.A[i + j, col]:
                            self.A[i] -= int(self.A[i, col] / self.A[i + j, col] - 1) * self.A[i + j]
                        else:
                            self.A[i] -= int(self.A[i, col] / self.A[i + j, col]) * self.A[i + j]
                    elif abs(self.A[i, col]) < abs(self.A[i + j, col]):
                        if i >= self.row_target:
                            if self.A[i + j, col] / self.A[i, col] == self.A[i + j, col] // self.A[i, col]:
                                self.A[i + j] -= int(self.A[i + j, col] / self.A[i, col] - 1) * self.A[i]
                            else:
                                self.A[i + j] -= int(self.A[i + j, col] / self.A[i, col]) * self.A[i]
                    i += 1
        # 再将右上角部分归零
        for col in range(self.row_target):
            flag = True
            while list(self.A[:col, col]).count(0) < col and flag:
                for r in range(col):
                    if abs(self.A[r, col] / self.A[col, col]) >= 1:
                        self.A[r] -= int(self.A[r, col] / self.A[col, col]) * self.A[col]
                    else:       # 比对角线处绝对值小就跳至下一列
                        flag = False
        # 将左下角彻底归零
        for col in range(self.row_target, self.cols):  # 遍历列
            for row in range(self.row_target, self.cols):   # 遍历行
                for j in range(row):      # 将下面每一行与上面全部行对比
                    s2 = sum(map(abs, self.A[j, self.row_target:]))
                    while 1:
                        if sum(map(abs, self.A[row, self.row_target:] - self.A[j, self.row_target:])) < s2:
                            self.A[j] -= self.A[row]
                            s2 = sum(map(abs, self.A[j, self.row_target:]))
                        elif sum(map(abs, self.A[row, self.row_target:] + self.A[j, self.row_target:])) < s2:
                            self.A[j] += self.A[row]
                            s2 = sum(map(abs, self.A[j, self.row_target:]))
                        else:
                            break

    # 处理参数矩阵，避免结果参数数值过大
    def solve_y(self):
        try:        # 同一实例下当已执行该方法则不会继续调用
            self.change_A_to_B()
        except:
            pass
        # 生成参数列表 -->形如[10, 2, t1, t2, t3]
        y_mat = np.linalg.solve(self.A[:self.row_target, :self.row_target].T, self.target)
        arg_ls = list(map(int, y_mat.ravel())) + [1 for i in range(self.cols - self.row_target)]
        B = self.A[:, self.row_target:]  # 切片出B矩阵
        X = []
        for col in range(self.cols):  # 列
            ls = [0]
            for row in range(self.cols):  # 行
                if row < self.row_target:
                    ls[0] = ls[0] + B[row, col] * arg_ls[row]
                else:
                    ls.append(B[row, col] * arg_ls[row])
            X += ls
        y_mat = np.array(X).reshape(self.cols, -1)
        left = sum(map(abs, y_mat[:, 0]))   # 尽量使常数项绝对值和最小
        for i in range(self.cols-self.row_target):      # 只遍历未知参数列
            while 1:
                if sum(map(abs, y_mat[:, 0] - y_mat[:, i+1])) < left:
                    y_mat[:, 0] -= y_mat[:, i+1]
                    left = sum(map(abs, y_mat[:, 0]))
                elif sum(map(abs, y_mat[:, 0] + y_mat[:, i+1])) < left:
                    y_mat[:, 0] += y_mat[:, i+1]
                    left = sum(map(abs, y_mat[:, 0]))
                else:
                    break
        # print(y_mat)
        return y_mat   # 返回结果矩阵，第一列常数项，后续为参数系数列

    def show_general_result(self):      # 通解形式
        y_mat = self.solve_y()
        X_ls = []
        for row in range(self.cols):
            ls = []
            if y_mat[row, 0] == 0:
                s = f'x_{row + 1} ='
            else:
                s = f'x_{row + 1} = {str(int(y_mat[row, 0]))}'
            for col in range(self.cols-self.row_target):
                if y_mat[row, col+1] == 1:
                    ls.append(f't_{col+1}')
                elif y_mat[row, col+1] == -1:
                    ls.append(f'-t_{col+1}')
                elif y_mat[row, col+1] != 0:
                    ls.append(str(int(y_mat[row, col+1])) + '·' + f't_{col+1}')
            for i in ls:
                if i[0] == '-':
                    s += ' - ' + i[1:]
                else:
                    s += ' + ' + i
            X_ls.append(s)
        return X_ls

    def show_specific_result(self, *args):      # 特解形式,传入参数需为各参数取值列表-->形如[[1, 5], [0, 10], [10, 100]]
        y_mat = self.solve_y()
        if args:
            pool = args[0]
            if len(pool) > self.cols - self.row_target:
                print(f'输入条件过多，当前条件下仅有{self.cols-self.row_target}个参数')
            else:
                value = pool[0]
                if isinstance(value, list):  # 范围输入分支
                    for val in self.rewrite_iter(pool):
                        mat = np.array([1] + list(val)).reshape(-1, 1)
                        y = y_mat.dot(mat)
                        yield list(map(int, y.ravel()))
                elif isinstance(value, tuple):  # 点输入分支
                    for val in pool:
                        mat = np.array([1] + list(val)).reshape(-1, 1)
                        y = y_mat.dot(mat)
                        yield list(map(int, y.ravel()))
                else:
                    print('以列表输入参数范围，元组输入参数取值点')
        else:
            # print('特解默认输出各参数0~5内的解')
            for value in itertools.product(range(10), repeat=self.cols-self.row_target):
                mat = [1] + list(value)
                y = y_mat.dot(np.array(mat).reshape(-1, 1))
                yield list(map(int, y.ravel()))

    def rewrite_iter(self, *args):
        pools = [[ls[0]+i for i in range(ls[1]-ls[0])] for ls in args[0]]
        result = [[]]
        for pool in pools:
            result = [x + [y] for x in result for y in pool]
        for prod in result:
            yield tuple(prod)


# 封装接口函数
# 若待解参数含小数，先整体扩大整10倍数再传入
def solve(coef, target, *args, show_specific=True):
    ind_equ = IndefinitEquation(coef, target)
    if args:
        result = np.array(list(ind_equ.show_specific_result(args))).reshape(-1, ind_equ.cols)
    else:
        if show_specific:
            result = np.array(list(ind_equ.show_specific_result())).reshape(-1, ind_equ.cols)
        else:
            result = np.array(ind_equ.show_general_result()).reshape(ind_equ.cols, -1)
    return result   # 不考虑可能出现大型方程组的内存消耗，增加查看便捷性将结果统一为矩阵输出
    # if args:
    #     result = np.array(list(ind_equ.show_specific_result(args))).reshape(-1, ind_equ.cols)
    # else:
    #     result = np.array(ind_equ.show_general_result()).reshape(ind_equ.cols, -1)
    # return result     # 特解返回生成器，通解返回列表


if __name__ == '__main__':
    np.set_printoptions(suppress=True)
    # a = [3, 5, 4, -11]
    # target = [10]
    # a = [3, 7, 16]
    # target = [40]
    # a = [49, -56, 14]
    # target = [35]
    # a = [4, 5, -1, 1, 6, 3, -9, -7, 5, -10, 5, 2, 1, -4, 7]
    # target = [3, 1, 2]
    a = [1, 1, 1, 3, 2, 1]
    target = [30, 50]
    # a = [111, -333, -259, 185, -370]
    # target = [148]
    r = solve(a, target)
    # r = solve(a, target, [1,4],[0,5])
    print(r)
    # for i in r:
    #     print(i)
    # print('\t' + ', '.join([f't_{i+1}' for i in range(solve.cols-solve.row_target)]) + ' ∈ Z')
    # # res = solve.show_specific_result((1,5), (2,3))
    # res = solve.show_specific_result([1,5], [2,4])
    # for i in res:
    #     print(i)
